package com.ai.zuochengyun.phase01.class06;

import com.ai.zuochengyun.phase01.class07.Node;
import com.ai.zuochengyun.phase01.class07.NodeRecord;

import java.util.HashMap;
import java.util.Map;

public class NodeHeap {

    private Node[] nodes;

    /**
     * 反向索引表
     */
    private Map<Node, Integer> heapIndexMap;

    private Map<Node, Integer> distanceMap;

    private int size;

    public NodeHeap(int capacity) {
        this.nodes = new Node[capacity];
        this.heapIndexMap = new HashMap<>();
        this.distanceMap = new HashMap<>();
        size = 0;
    }

    public void addOrUpdateOrIgnore(Node node, int distance) {
        if (!isEntered(node)) {
            // 执行插入操作
            nodes[size] = node;
            distanceMap.put(node, distance);
            heapIndexMap.put(node, size);
            // 小的节点上浮
            heapInsert(size);
            size++;
        }
        if (isInHeap(node)) {
            // 执行更新操作
            distanceMap.put(node, Math.min(distanceMap.get(node), distance));
            // 小的节点上浮
            heapInsert(heapIndexMap.get(node));
        }
    }

    public NodeRecord pop() {
        NodeRecord nodeRecord = new NodeRecord(nodes[0], distanceMap.get(nodes[0]));

        int tail = size -1;

        // 弹出位置的节点和尾节点交换位置，然后尾节点需要进行heapify
        swap(0, tail);
        heapIndexMap.put(nodes[tail], -1);
        distanceMap.remove(nodes[tail]);
        nodes[tail] = null;

        // 大的节点下沉
        heapify(0, size);
        size--;

        // 返回头节点
        return nodeRecord;
    }

    public boolean isEmpty() {
        return size == 0;
    }

    /**
     * 有没有进来过
     * @param node
     * @return
     */
    public boolean isEntered(Node node) {
        return heapIndexMap.containsKey(node);
    }

    /**
     * 在不在堆上
     * @param node
     * @return
     */
    public boolean isInHeap(Node node) {
        return heapIndexMap.get(node) != -1;
    }

    /**
     * 小的值往上移动
     * @param index
     */
    public void heapInsert(int index) {
        while (distanceMap.get(nodes[index]) < distanceMap.get(nodes[(index -1)/2])) {
            swap((index -1)/2, index);
            index = (index -1)/2;
        }
    }

    /**
     * 大的值往下移动
     * @param index
     * @param size
     */
    public void heapify(int index, int size) {
        int left = index*2+1;
        while (left < size) {
            // 右侧有元素，确定哪个子节点更小
            int smallest = left+1 < size && distanceMap.get(nodes[left+1]) > distanceMap.get(nodes[left]) ? left+1 : left;
            // 父节点和最小子节点的值进行比较
            smallest = distanceMap.get(nodes[smallest]) < distanceMap.get(nodes[index]) ? smallest : index;
            if (smallest == index) {
                return;
            }

            swap(smallest, index);
            index = smallest;
            left = index*2+1;
        }
    }

    public void swap(int smallest, int index) {
        // 反向索引表要同步交换
        heapIndexMap.put(nodes[smallest], index);
        heapIndexMap.put(nodes[index], smallest);
        Node temp = nodes[smallest];
        nodes[smallest] = nodes[index];
        nodes[index] = temp;
    }
}
